import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat


def make_dataset():
    orl = {'nr_class': 40, 'nr_train': 320, 'nr_test': 80, 'dim': 644}
    veh = {'nr_class': 4, 'nr_train': 677, 'nr_test': 169, 'dim': 18}
    ORL_DATA = np.array(loadmat('data/ORLData_25_400_645.mat')['X']).T
    VEH_DATA = np.array(loadmat('data/vehicle_846_19.mat')['X']).T
    orl['train'], orl['test'] = ORL_DATA[:-1, :8], ORL_DATA[:-1, 8:10]
    orl['train_label'], orl['test_label'] = ORL_DATA[-1, :8], ORL_DATA[-1, 8:10]
    veh['train'], veh['test'] = VEH_DATA[:-1, :677], VEH_DATA[:-1, 677:]
    veh['train_label'], veh['test_label'] = VEH_DATA[-1, :677], VEH_DATA[-1, 677:]
    for i in range(1, 40):
        orl['train'] = np.c_[orl['train'], ORL_DATA[:-1, i * 10:i * 10 + 8]]
        orl['test'] = np.c_[orl['test'], ORL_DATA[:-1, i * 10 + 8:i * 10 + 10]]
        orl['train_label'] = np.r_[orl['train_label'], ORL_DATA[-1, i * 10:i * 10 + 8]]
        orl['test_label'] = np.r_[orl['test_label'], ORL_DATA[-1, i * 10 + 8:i * 10 + 10]]
    return orl, veh


def PCA(dataset, low_dim):
    X_ori = dataset['train']
    means = np.mean(X_ori, 1).reshape(dataset['dim'], 1)
    X = X_ori - means
    cov = X.dot(X.T)
    eig_val, eig_vec = np.linalg.eig(cov)
    W = eig_vec[:, np.argsort(eig_val)[-low_dim:]]
    return W.T.dot(X_ori), W.T.dot(dataset['test'])


def KNN_1(dataset, low_dim_train, low_dim_test):
    correct = 0
    for i in range(dataset['nr_test']):
        best_index, current_best = -1, float('inf')
        for j in range(dataset['nr_train']):
            dist = np.linalg.norm(low_dim_train[:, j] - low_dim_test[:, i])
            if dist < current_best:
                current_best = dist
                best_index = j
        if dataset['train_label'][best_index] == dataset['test_label'][i]:
            correct += 1
    return correct / dataset['nr_test']


def LDA(dataset, low_dim):
    S_w = np.zeros((dataset['dim'], dataset['dim']))
    S_b = np.zeros((dataset['dim'], dataset['dim']))
    all_mean = np.mean(dataset['train'], 1).reshape(dataset['dim'], 1)
    for i in range(dataset['nr_class']):
        mean, cnt, x_tmp = np.zeros((dataset['dim'], 1)), 0, None
        for j in range(dataset['nr_train']):
            if dataset['train_label'][j] == i + 1:
                sample = dataset['train'][:, j].reshape(dataset['dim'], 1)
                mean += sample
                cnt += 1
                if x_tmp is None:
                    x_tmp = sample
                else:
                    x_tmp = np.c_[x_tmp, sample]
        mean /= cnt
        x_tmp = x_tmp - mean
        mean -= all_mean
        S_w += x_tmp.dot(x_tmp.T)
        S_b += cnt * mean.dot(mean.T)
    eig_val, eig_vec = np.linalg.eigh(np.linalg.inv(S_w).dot(S_b))
    W = eig_vec[:, np.argsort(eig_val)[-low_dim:]]
    return W.T.dot(dataset['train']), W.T.dot(dataset['test'])


def PCA_KNN_ORL(orl):
    # PCA + KNN with ORL
    dim, acc = [], []
    for ld in range(5, orl['dim'], 5):
        dim.append(ld)
        l_train, l_test = PCA(orl, ld)
        acc.append(KNN_1(orl, l_train, l_test))
    plt.plot(dim, acc)
    plt.title('PCA + KNN with ORL')
    plt.xlabel('low_dim')
    plt.ylabel('test accuracy')
    plt.savefig('D:\\Document\\2021\\课程\\模式识别\\Homework\\HW5\\latex\\picture\\PCA_KNN_ORL.pdf')


def PCA_KNN_VEH(veh):
    # PCA + KNN with VEH
    dim, acc = [], []
    for ld in range(1, veh['dim']):
        dim.append(ld)
        l_train, l_test = PCA(veh, ld)
        acc.append(KNN_1(veh, l_train, l_test))
    plt.plot(dim, acc)
    plt.title('PCA + KNN with VEH')
    plt.xlabel('low_dim')
    plt.ylabel('test accuracy')
    plt.savefig('D:\\Document\\2021\\课程\\模式识别\\Homework\\HW5\\latex\\picture\\PCA_KNN_VEH.pdf')


def LDA_KNN_ORL(orl):
    # LDA + KNN with ORL
    dim, acc = [], []
    for ld in range(5, orl['dim'], 5):
        dim.append(ld)
        l_train, l_test = LDA(orl, ld)
        acc.append(KNN_1(orl, l_train, l_test))
    plt.plot(dim, acc)
    plt.title('LDA + KNN with ORL')
    plt.xlabel('low_dim')
    plt.ylabel('test accuracy')
    plt.savefig('D:\\Document\\2021\\课程\\模式识别\\Homework\\HW5\\latex\\picture\\LDA_KNN_ORL.pdf')


def LDA_KNN_VEH(veh):
    # PCA + KNN with VEH
    dim, acc = [], []
    for ld in range(1, veh['dim']):
        dim.append(ld)
        l_train, l_test = LDA(veh, ld)
        acc.append(KNN_1(veh, l_train, l_test))
    plt.plot(dim, acc)
    plt.title('LDA + KNN with VEH')
    plt.xlabel('low_dim')
    plt.ylabel('test accuracy')
    plt.savefig('D:\\Document\\2021\\课程\\模式识别\\Homework\\HW5\\latex\\picture\\LDA_KNN_VEH.pdf')


if __name__ == '__main__':
    ORL, VEH = make_dataset()
    # PCA_KNN_ORL(ORL)
    # PCA_KNN_VEH(VEH)
    # LDA_KNN_ORL(ORL)
    LDA_KNN_VEH(VEH)
